


import chisel3._
import chisel3.util._
import ID._

class Mem extends Module {
  val io = IO(new Bundle {
    val pc       = Input(UInt(32.W))    
    val op_type  = Input(UInt(6.W))
    val rs1      = Input(UInt(64.W))
    val rs2      = Input(UInt(64.W))
    val src2     = Input(UInt(64.W))
    val fuType   = Input(UInt(3.W))   
    val irq      = Input(Bool())
    val mem_en   = Output(Bool())
    val rd_data  = Output(UInt(64.W))
    val busy     = Output(Bool())
    val dmem_hs  = Output(Bool())
    val resp_success = Output(Bool())
    val dmem     = new CacheBusIO
  })

//  val ls_axi_id = 2.U(AxiParameters.AxiIdWidth.W)   // id = 2 for load/store
  val s_idle :: s_req :: s_wait_r :: s_wait_w :: s_complete :: Nil = Enum(5)
  val state = RegInit(s_idle)

  val req = io.dmem.req
  val resp = io.dmem.resp

  val addr = (io.rs1 + io.rs2)(31,0)
  val addr_offset = addr(2, 0)
  val wdata = io.src2
  val reg_addr = RegInit(0.U(32.W))
  val reg_addr_offset = reg_addr(2, 0)
  val reg_wdata = RegInit(0.U(64.W))

  val mem_en  = io.fuType === FuLsu 
  val mem_ren = io.fuType === FuLsu && (io.op_type(3) === 0.U) 
  val mem_wen = io.fuType === FuLsu && io.op_type(3) 
  io.mem_en := mem_en

  val mmio = (addr(31) === 0.U) && (io.pc =/= "h00000000".U)

  val mask = MuxLookup(reg_addr_offset, 0.U, Array(
    "b000".U -> "b11111111".U,
    "b001".U -> "b11111110".U,
    "b010".U -> "b11111100".U,
    "b011".U -> "b11111000".U,
    "b100".U -> "b11110000".U,
    "b101".U -> "b11100000".U,
    "b110".U -> "b11000000".U,
    "b111".U -> "b10000000".U
  ))
  val wmask = MuxLookup(io.op_type, 0.U, Array(
    LsuSb  -> "b00000001".U,
    LsuSh  -> "b00000011".U,
    LsuSw  -> "b00001111".U,
    LsuSd  -> "b11111111".U
  ))
 
//  req.bits.id := ls_axi_id
  req.bits.addr := Mux(mmio, addr, Cat(reg_addr(31, 3), Fill(3, 0.U)))
  req.bits.ren := mem_ren
  req.bits.wdata := (reg_wdata << (reg_addr_offset << 3))(63, 0)
  req.bits.wmask := mask & ((wmask << reg_addr_offset)(7, 0))
  req.bits.wen := mem_wen
  req.bits.size := Mux(mmio, io.op_type, "b11".U)
  
  req.valid := (state === s_req) && (mem_ren || mem_wen) && !io.irq

  resp.ready := (state === s_wait_r) || (state === s_wait_w)

  val ram_rdata = RegInit(UInt(64.W), 0.U)
  // val resp_r_success = resp.fire() && resp.bits.rlast &&
  //                      (resp.bits.id === ls_axi_id)
  // val resp_w_success = resp.fire() && resp.bits.wresp &&
  //                      (resp.bits.id === ls_axi_id)
  io.resp_success := resp.fire()
  io.dmem_hs := resp.ready && resp.valid
  
  switch (state) {
    is (s_idle) {
      when (mem_en & !io.irq) {
        state := s_req
        reg_addr := addr
        reg_wdata := wdata
      }
    }
    is (s_req) {
      when (io.irq) {
        state := s_idle      
      } .elsewhen (mem_ren && req.fire()) {
        state := s_wait_r
      } .elsewhen (mem_wen && req.fire()) {
        state := s_wait_w
      }
    }
    is (s_wait_r) {
      when (io.resp_success) {
        ram_rdata := resp.bits.rdata >> (reg_addr_offset << 3)
        state := s_complete
      }
    }
    is (s_wait_w) {
      when (io.resp_success) {
        state := s_complete
      }
    }
    is (s_complete) {
      state := s_idle
      reg_addr := 0.U
    }
  }

  val load_out = Wire(UInt(64.W))

  load_out :=  MuxLookup(io.op_type, 0.U, Array(
      LsuLb   -> Cat(Fill(56, ram_rdata(7)), ram_rdata(7, 0)),      
      LsuLh   -> Cat(Fill(48, ram_rdata(15)), ram_rdata(15, 0)),
      LsuLw   -> Cat(Fill(32, ram_rdata(31)), ram_rdata(31, 0)),
      LsuLd   -> ram_rdata,
      LsuLbu  -> Cat(Fill(56, 0.U), ram_rdata(7, 0)),
      LsuLhu  -> Cat(Fill(48, 0.U), ram_rdata(15, 0)),
      LsuLwu  -> Cat(Fill(32, 0.U), ram_rdata(31, 0)),  
    ))

  io.rd_data := Mux(state === s_complete, load_out, 0.U)
  io.busy := ((state === s_idle) && mem_en) || (state === s_req) ||
             (state === s_wait_r) || (state === s_wait_w)

}
